import json
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel  # 用于加载 LoRA adapter
import re  # 用于从模型输出中提取选项

# --- 配置参数 ---
base_model_path = '/hy-tmp/models/Qwen-0.6B/'  # 您的基础模型路径 (请使用绝对路径)
adapter_path = '/hy-tmp/SceneryMC/sft/final_model_peft/'  # 您的 LoRA adapter 路径 (请使用绝对路径)
eval_data_path = '/hy-tmp/SceneryMC/sft/data/eval.jsonl'  # 您的 eval.jsonl 文件路径

# 模型生成参数
max_new_tokens = 1000  # 模型生成答案的最大 token 数，选项字母通常很短
temperature = 0.1  # 控制生成随机性，较低的值使输出更确定
# top_p = 0.9         # 可以根据需要调整

device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Using device: {device}")

# --- 1. 加载模型和分词器 ---
print(f"Loading base model from: {base_model_path}")
tokenizer = AutoTokenizer.from_pretrained(
    base_model_path,
    trust_remote_code=True,
    padding_side='left'  # 对于生成任务，通常左填充
)
if tokenizer.pad_token is None:
    if tokenizer.eos_token is not None:
        tokenizer.pad_token = tokenizer.eos_token
        print(f"Set tokenizer.pad_token to eos_token: {tokenizer.eos_token}")
    else:
        # 对于Qwen3，不太可能发生
        tokenizer.add_special_tokens({'pad_token': '[PAD]'})
        print(f"Added new pad_token: {tokenizer.pad_token}")

# 加载基础模型
model = AutoModelForCausalLM.from_pretrained(
    base_model_path,
    trust_remote_code=True,
    torch_dtype=torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16,  # 使用半精度以节省显存和加速
    device_map="auto"  # 自动分配到可用设备
)
print(f"Base model loaded. Vocab size: {model.config.vocab_size}")

# 加载 LoRA adapter
# print(f"Loading LoRA adapter from: {adapter_path}")
# try:
#     model = PeftModel.from_pretrained(model, adapter_path, is_trainable=False)
#     model = model.merge_and_unload()  # 可选：如果想合并权重并卸载LoRA层以加快推理速度 (但会增加显存占用)
#     # 如果不合并，直接用 PeftModel 也可以推理
#     print("LoRA adapter loaded and merged (if applicable).")
# except Exception as e:
#     print(f"Error loading LoRA adapter: {e}")
#     print("Proceeding with base model only (if adapter loading failed).")

model.eval()  # 设置为评估模式


# --- 2. 读取数据集 ---
def load_eval_data(file_path):
    data = []
    with open(file_path, 'r', encoding='utf-8') as f:
        for line in f:
            if line.strip():
                try:
                    data.append(json.loads(line))
                except json.JSONDecodeError:
                    print(f"Skipping invalid JSON line: {line.strip()}")
    return data


eval_dataset = load_eval_data(eval_data_path)
print(f"Loaded {len(eval_dataset)} samples from {eval_data_path}")


# --- 3. 构建 Prompt 的函数 ---
def build_prompt(item):
    question = item['question']
    options_list = item['options']
    topic = item.get('topic', 'Unknown Topic')

    # 格式化选项 A) B) C) D)
    formatted_options = []
    for i, opt_text in enumerate(options_list):
        letter = chr(ord('A') + i)
        formatted_options.append(f"{letter}) {opt_text}")
    options_str = "\n".join(formatted_options)

    # 构建完整的对话 prompt，这应该与您 SFT 训练时的格式类似或兼容
    # 这里使用一个示例 prompt，请根据您的实际情况调整
    messages = [
        {"role": "system",
         "content": f"你是一个专业的AI助手，领域是{topic}。请仔细阅读下面的单项选择题。你的任务是只输出一个代表正确答案的选项字母（例如：A, B, C, 或 D）。不要包含任何其他文字、解释、思考过程或特殊标签，。"},
        {"role": "user",
         "content": f"题目：{question}\n\n选项：\n{options_str}\n\n请严格按照要求，仅输出正确答案的选项字母："}
    ]

    # 使用 tokenizer 的 apply_chat_template 来格式化对话
    # add_generation_prompt=True 会在末尾添加助手开始的标记，如 <|im_start|>assistant\n
    try:
        prompt_text = tokenizer.apply_chat_template(
            messages,
            tokenize=False,
            add_generation_prompt=True
        )
    except Exception as e:
        print(f"Error applying chat template: {e}")
        # 备用简单拼接方式 (如果模板应用失败)
        prompt_text = f"<|im_start|>system\n你是一个专业的AI助手，领域是{topic}。请仔细阅读下面的单项选择题，然后只输出你认为正确的选项字母（例如：A, B, C, 或 D）。<|im_end|>\n"
        prompt_text += f"<|im_start|>user\n题目：{question}\n\n选项：\n{options_str}\n\n正确答案的选项字母是：<|im_end|>\n<|im_start|>assistant\n"

    return prompt_text, len(options_list)


# --- 4. 从模型回答中提取选项的函数 ---

def extract_option(model_output_answer_part, num_options):
    """
    从模型生成的答案部分文本中提取选项字母 (A, B, C, D, ...)。
    会尝试处理 <think> 标签等常见干扰，并更鲁棒地查找选项。
    """
    text_to_search = model_output_answer_part
    print(f"    DEBUG: Text to search in extract_option: '{text_to_search[:100]}'")  # 打印用于搜索的文本

    # 步骤 1: 尝试移除已知的思考标签或前缀
    # Qwen 系列可能使用 <think> ... </think> 或类似结构
    # 或者直接以 <think> 开头
    think_pattern = re.compile(r"<think>(.*?)</think>", re.DOTALL)
    match_think_block = think_pattern.search(text_to_search)
    if match_think_block:
        # 如果找到完整的 <think>...</think> 块，看块后面是否有内容
        remaining_text_after_think = text_to_search[match_think_block.end():].strip()
        if remaining_text_after_think:
            text_to_search = remaining_text_after_think
            print(f"    DEBUG: Removed <think> block, remaining: '{text_to_search[:50]}'")
        else:  # 如果 </think> 后面没内容了，就用 </think> 之前的内容
            text_to_search = match_think_block.group(1).strip()  # 取 <think> 和 </think> 之间的内容
            print(f"    DEBUG: Using content within <think> block: '{text_to_search[:50]}'")

    # 另一种可能是只以 <think> 开头，没有结束标签，或者我们只关心它之后的内容
    if text_to_search.lower().startswith("<think>"):
        # 尝试找到 <think> 标签后的第一个换行符或一个有意义的分隔符
        # 或者简单地去掉 <think>
        think_prefix_match = re.match(r"<think>\s*", text_to_search, re.IGNORECASE)
        if think_prefix_match:
            text_to_search = text_to_search[think_prefix_match.end():].strip()
            print(f"    DEBUG: Removed <think> prefix, remaining: '{text_to_search[:50]}'")

    # 也可以移除其他常见的前缀，例如 "答案是：", "好的，", "The answer is " 等
    common_prefixes = [
        r"答案是\s*[:：]?\s*",
        r"好的，\s*(?:答案是\s*[:：]?\s*)?",
        r"the correct answer is\s*[:：]?\s*",
        r"option\s*[:：]?\s*",
        r"predicted option\s*[:：]?\s*",
        r"assistant\s*[:：]?\s*",  # 如果 apply_chat_template 后仍然有残留
    ]
    for prefix_regex in common_prefixes:
        match_prefix = re.match(prefix_regex, text_to_search, re.IGNORECASE)
        if match_prefix:
            text_to_search = text_to_search[match_prefix.end():].strip()
            print(f"    DEBUG: Removed prefix '{match_prefix.group(0)}', remaining: '{text_to_search[:50]}'")
            break  # 通常移除一个就够了

    # 步骤 2: 精确匹配独立的选项字母 (例如 "A", "B.", "(C)", " D ")
    valid_options_letters = [chr(ord('A') + i) for i in range(num_options)]
    for letter in valid_options_letters:
        # 匹配以字母开头，后面可能跟点、括号，或者被空格包围
        # \b 是单词边界，确保匹配的是独立的字母而不是单词的一部分
        # (?: ... ) 是非捕获组
        patterns_for_letter = [
            r"^\s*" + re.escape(letter) + r"\s*(?:[:：\.,\)\s]|$)",  # 开头是字母，后面是分隔符或结尾
            r"\b" + re.escape(letter) + r"\b"  # 独立的字母
        ]
        for pattern_str in patterns_for_letter:
            match = re.match(pattern_str, text_to_search, re.IGNORECASE)  # 从开头匹配
            if match:
                print(f"    DEBUG: Matched pattern '{pattern_str}' for letter '{letter}'")
                return letter.upper()  # 确保返回大写

    # 步骤 3: 如果上面没匹配到，再尝试查找文本中出现的第一个有效选项字母
    # 这比较宽松，可能会误判
    first_char_match = re.search(r"([A-Z])", text_to_search)  # 查找任何大写字母
    if first_char_match:
        potential_letter = first_char_match.group(1).upper()
        if potential_letter in valid_options_letters:
            print(f"    DEBUG: Fallback: Found first valid capital letter '{potential_letter}'")
            return potential_letter

    print(f"    Warning: Could not extract a valid option letter from processed text: '{text_to_search[:50]}...'")
    return None


# --- 5. 进行评测 ---
eval_batch_size = 8  # 或者 16, 32，根据您的显存调整

correct_predictions = 0
total_predictions = 0
num_batches = (len(eval_dataset) + eval_batch_size - 1) // eval_batch_size

for batch_idx in range(num_batches):
    print(f"\nProcessing batch {batch_idx + 1}/{num_batches}...")
    batch_start_idx = batch_idx * eval_batch_size
    batch_end_idx = min((batch_idx + 1) * eval_batch_size, len(eval_dataset))
    current_batch_data = eval_dataset[batch_start_idx:batch_end_idx]

    if not current_batch_data:
        continue

    prompts_and_options_info = [build_prompt(item) for item in current_batch_data]
    batch_prompts = [p_info[0] for p_info in prompts_and_options_info]
    batch_num_options = [p_info[1] for p_info in prompts_and_options_info]
    batch_true_answers = [item['answer'].upper() for item in current_batch_data]

    # 对整个批次的 prompts 进行 tokenize 和 padding
    # padding=True 和 truncation=True (如果需要) 会在这里确保所有序列长度一致
    inputs = tokenizer(
        batch_prompts,
        return_tensors="pt",
        padding="longest",  # 填充到批次中最长的序列
        truncation=True,  # 如果有超长的prompt，进行截断
        max_length=tokenizer.model_max_length if hasattr(tokenizer, 'model_max_length') else 2048  # 或您设定的最大长度
    ).to(device)

    # 模型生成
    with torch.no_grad():
        try:
            outputs = model.generate(
                **inputs,  # inputs 包含了 input_ids, attention_mask
                max_new_tokens=max_new_tokens,
                temperature=temperature,
                do_sample=True if temperature > 0 else False,
                pad_token_id=tokenizer.pad_token_id,
                eos_token_id=tokenizer.eos_token_id
            )

            # 解码生成的 token IDs
            # outputs 的形状是 (batch_size, sequence_length)
            # generated_texts_full = tokenizer.batch_decode(outputs, skip_special_tokens=True)

            # 只解码新生成的部分
            generated_texts_answer_part = []
            for i in range(outputs.shape[0]):  # 遍历批次中的每个样本
                # input_ids 的长度可能因padding而不同，需要取每个样本的实际prompt长度
                # 或者，更简单的是，如果tokenizer的padding_side是'left'，那么真实内容在右边
                # inputs['input_ids'][i] 是已经padding过的prompt
                # 我们需要找到每个prompt中非padding部分的结束位置
                # 一个简单的方法是假设 generate() 返回的 outputs 包含了输入的 prompt
                # 并且新生成的 token 在其后

                # 如果 tokenizer.padding_side == 'left'
                # prompt_len = (inputs['input_ids'][i] != tokenizer.pad_token_id).sum().item()
                # generated_ids_for_sample = outputs[i][prompt_len:]

                # 或者，更通用的方式，generate 通常返回完整的序列（prompt + generation）
                # 而我们知道 inputs['input_ids'].shape[1] 是这个批次中 (可能被截断的) prompt 的最大长度 (因为 padding='longest')
                # 但每个原始 prompt 的长度可能不同。
                # 最简单的方法是解码整个序列，然后从解码后的文本中分离答案。
                # 或者，如果可以确定 inputs['input_ids'] 就是原始的、未被截断太多的输入
                # 且 padding_side='left'
                # 那么新生成的token就是 outputs[i][inputs['input_ids'].shape[1]:]

                # 一个常用的方法是直接解码 outputs，然后通过对比或特殊标记分离答案
                # 为了简单起见，我们假设 outputs[i] 的开头是 prompt，后面是答案
                # 如果使用了 apply_chat_template，且有 assistant 标记，可以从那里开始

                # 假设 inputs['input_ids'].shape[1] 是经过padding/truncation后的prompt长度
                prompt_length_in_batch = inputs['input_ids'].shape[1]
                generated_ids_for_sample = outputs[i][prompt_length_in_batch:]
                decoded_answer_part = tokenizer.decode(generated_ids_for_sample, skip_special_tokens=True)
                generated_texts_answer_part.append(decoded_answer_part)

                print(f"  Sample {batch_start_idx + i + 1}:")
                # print(f"    Full model output (decoded): {tokenizer.decode(outputs[i], skip_special_tokens=True)}") # 完整输出
                print(f"    Model answer part: '{decoded_answer_part}'")


        except Exception as e:
            print(f"  Error during model generation for batch {batch_idx + 1}: {e}")
            # 可以选择跳过整个批次或逐个处理错误
            total_predictions += len(current_batch_data)  # 假设整个批次都失败了
            continue

    # 逐个提取和比较答案
    for i in range(len(current_batch_data)):
        predicted_letter = extract_option(generated_texts_answer_part[i], batch_num_options[i])
        true_answer_letter = batch_true_answers[i]

        if predicted_letter:
            print(f"    True answer: {true_answer_letter}, Predicted option: {predicted_letter}")
            if predicted_letter == true_answer_letter:
                correct_predictions += 1
        else:
            print(f"    True answer: {true_answer_letter}, Prediction extraction failed.")

        total_predictions += 1

    # (可选) 少量样本后中断，用于测试
    # if total_predictions >= 5:
    #     break

# --- 6. 计算并打印结果 ---
if total_predictions > 0:
    accuracy = (correct_predictions / total_predictions) * 100
    print(f"\n--- Evaluation Finished ---")
    print(f"Total samples evaluated: {total_predictions}")
    print(f"Correct predictions: {correct_predictions}")
    print(f"Accuracy: {accuracy:.2f}%")
else:
    print("\n--- Evaluation Finished ---")
    print("No samples were evaluated.")